import torch
import mindspore as ms
import numpy as np
# 转化权重prompt
# 针对module.feature.conv数字开头的，比如module.feature.conv0.bn.weight，把conv0这种之后的：把conv.weight变成conv.weight不变，把bn.weight变成bn.gamma，bn.bias变成bn.beta，bn.running_mean变成bn.moving_mean，bn.running_var变成bn.moving_variance。module.feature.conv0.bn.num_batches_tracked不变
# 针对module.cost_regularization.conv数字开头的，比如module.cost_regularization.conv0.conv.weight，把conv0这种之后的：把conv.weight变成conv.weight不变，bn.weight变成bn.gamma，bn.bias变成bn.beta，bn.running_mean变成bn.moving_mean，bn.running_var变成bn.moving_variance，.bn.num_batches_tracked: ()不要

def map_name(name: str):
    # 删除前缀
    if name.startswith("module."):
        ms_name = name[len("module."):]
    else:
        ms_name = name

    # feature 分支 BN 参数
    if "feature" in ms_name and ".bn." in ms_name:
        ms_name = ms_name.replace(".bn.weight", ".bn.gamma")
        ms_name = ms_name.replace(".bn.bias", ".bn.beta")
        ms_name = ms_name.replace(".bn.running_mean", ".bn.moving_mean")
        ms_name = ms_name.replace(".bn.running_var", ".bn.moving_variance")
        # num_batches_tracked 保留
        return ms_name

    # cost_regularization 普通 BN 参数
    if "cost_regularization" in ms_name and ".bn." in ms_name:
        ms_name = ms_name.replace(".bn.weight", ".bn.bn2d.gamma")
        ms_name = ms_name.replace(".bn.bias", ".bn.bn2d.beta")
        ms_name = ms_name.replace(".bn.running_mean", ".bn.bn2d.moving_mean")
        ms_name = ms_name.replace(".bn.running_var", ".bn.bn2d.moving_variance")
        # num_batches_tracked 保留
        return ms_name

    # cost_regularization convX.Y 特殊分支
    # if "cost_regularization" in ms_name and (".0.weight" in ms_name):
    #     return ms_name.replace(".0.weight", ".0.weight")
    # if "cost_regularization" in ms_name and (".1.weight" in ms_name):
    #     return ms_name.replace(".1.weight", ".1.bn2d.gamma")
    # if "cost_regularization" in ms_name and (".1.bias" in ms_name):
    #     return ms_name.replace(".1.bias", ".1.bn2d.beta")
    # if "cost_regularization" in ms_name and (".1.running_mean" in ms_name):
    #     return ms_name.replace(".1.running_mean", ".1.bn2d.moving_mean")
    # if "cost_regularization" in ms_name and (".1.running_var" in ms_name):
    #     return ms_name.replace(".1.running_var", ".1.bn2d.moving_variance")
    # # .1.num_batches_tracked 保留原样
    # if "cost_regularization" in ms_name and (".1.num_batches_tracked" in ms_name):
    #     return ms_name

    return ms_name

def torch_to_ms(torch_ckpt_path, ms_ckpt_path):
    torch_ckpt = torch.load(torch_ckpt_path, map_location="cpu")
    if "model" in torch_ckpt:
        state_dict = torch_ckpt["model"]
    elif "state_dict" in torch_ckpt:
        state_dict = torch_ckpt["state_dict"]
    else:
        state_dict = torch_ckpt

    new_params = []
    for name, param in state_dict.items():
        # 跳过 num_batches_tracked 结尾的参数
        if name.endswith('num_batches_tracked'):
            continue
        ms_name = map_name(name)
        if ms_name is None:
            continue
        data = param.cpu().numpy().astype(np.float32)
        new_params.append({"name": ms_name, "data": ms.Tensor(data)})

    ms.save_checkpoint(new_params, ms_ckpt_path)
    print(f"✅ 已保存 MindSpore ckpt: {ms_ckpt_path}")

if __name__ == "__main__":
    torch_to_ms("checkpoints/casmvsnet.ckpt", "model_000000ms.ckpt")
